#pragma once
#include <functional>
#include <fstream>

// BOOST
#include <boost/mpi.hpp>
#include <boost/program_options.hpp>
// ARMADILLO
#include <armadillo>
// HELPER
#include "../helper/helper.h"
#include "../helper/mpi_helper.h"

using namespace std;
using namespace arma;
namespace mpi = boost::mpi;
namespace po = boost::program_options;

namespace factor {

  struct Parameter {      
    int m, n, k;
    int verbose;
    int maxiter;
    double epsilon, tol;
    string output;
    double alpha;
      
    void Print() {
      MSG("m=" << m << " n=" << n << " k=" << k 
	  << " maxiter=" << maxiter << " epsilon=" << epsilon
	  << " tol=" << tol);
    }
  }; // class Parameter

  struct Iteration {
    int iter;
    double normA, normE, relE;
    posix::time_duration duration;
    double spW, spH;
    double mse;

    Iteration(int i, double A, double E, double rE, 
	      posix::time_duration d, double sW, double sH, double mse_)
    : iter(i), normA(A), normE(E), relE(rE), duration(d), spW(sW), spH(sH), mse(mse_) {}

    void Print() {
      MSG("iter=" << iter << " |WH'-A|/|A|=" << relE << " rmse=" << mse 
	  << " d=" << SECONDS(duration) << " spW=" << spW << " spH=" << spH);
    }

    void Output(ostream& out) {
      out << iter << ",\t" << SECONDS(duration) << ",\t" 
	  << relE << ",\t" << mse << ",\t"
	  << spW << ",\t" << spH
	  << endl;
    }    
  }; // class Iteration

  template<typename MAT> class MultNMF {
  public:
    mpi::communicator& world;
    po::variables_map& vm;
    Parameter par;

    posix::ptime startTime;
    vector<Iteration> progress;
    mat W_big, H_big, W, H;
    int mr, nc; // number of rows of local W, H
    double normAlocal, normA;
    double normElocal, normE; // progress
    double relE, oldE;
    bool stop;
    
    MultNMF(mpi::communicator& world_, po::variables_map& vm_) 
      : world(world_), vm(vm_) {}
    
    void NMF(MAT& A);  // single process
    void NMF(MAT& Arows, MAT& Acols);  // multiple processes

    void Initialize();
    void ParInitialize();
    void ComputeNormA(MAT& Arows);
    void ComputeNormE(MAT& Arows);
    void Combine(mat& X, mat& X_big);

    void ProcessParameter();

    bool StopCriteriaMet();
    bool PrintCondition(int iter);
    void OutputProgress(Iteration& info);
  };

  template<typename MAT> void MultNMF<MAT>::NMF(MAT& A) {
    ARMA<MAT>::getSize(A, par.m, par.n);
    ProcessParameter();
    
    W = randu<mat>(par.m, par.k);
    H = randu<mat>(par.n, par.k);
    
    FORE(iter, 0, par.maxiter) {
      H = H % (trans(A)*W) / (H*(trans(W)*W) + par.epsilon);
      W = W % (A*H)        / (W*(trans(H)*H) + par.epsilon);
      if (par.verbose == 2) {
	MSG("iter=" << iter << " |WH'-A|=" << norm(A-W*trans(H),"fro"));
      }
    }
  }

  template<typename MAT> void MultNMF<MAT>::ProcessParameter() {
    REQUIRE_PARAM("k");
    
    par.k = vm["k"].as<int>();
    par.verbose = vm["verbose"].as<int>();
    par.maxiter = vm["iter"].as<int>();
    par.epsilon = vm["epsilon"].as<double>();
    par.tol = vm["tol"].as<double>();
    if ( vm["output"].as<string>() == "NONE" )
      par.output = "";
    else {
      stringstream ss;
      ss << vm["outputDir"].as<string>() << "/" << vm["output"].as<string>() << "_" << SIZE(world);
      par.output = ss.str();
    }
    par.alpha = vm["alpha"].as<double>();

    if (par.verbose && world.rank() == 0) par.Print();
    if (!par.output.empty()) { ofstream f(par.output.c_str(), ios::out); }
  }

  template<typename MAT> void MultNMF<MAT>::NMF(MAT& Arows, MAT& Acols) {
    this->startTime = CLOCK;
    this->progress.clear();

    ARMA<MAT>::getSize(Arows, mr, this->par.n);
    ARMA<MAT>::getSize(Acols, this->par.m, nc);
    this->ProcessParameter();
    
    ParInitialize();
    Combine(W, W_big);
    ComputeNormA(Arows);

    stop = false;
    relE = 1e-6;
    FORE(iter, 0, this->par.maxiter) {
      H = H % (trans(Acols)*W_big) / (H*(trans(W_big)*W_big) + this->par.epsilon);
      Combine(H, H_big);

      W = W % (Arows*H_big)        / (W*(trans(H_big)*H_big) + this->par.epsilon);
      Combine(W, W_big);

      ComputeNormE(Acols);
      oldE = relE;
      relE = normE / normA;
      if (MPI_IS_MASTER) stop = StopCriteriaMet();
      broadcast(this->world, stop, 0);

      if (MPI_IS_MASTER) {
	double spW = ARMA<mat>::sparsity(W_big), spH = ARMA<mat>::sparsity(H_big);
	Iteration info(iter, normA, normE, relE, CLOCK - this->startTime, spW, spH, sqrt(normE*normE / par.m / par.n));
	this->progress.push_back(info);

	if (PrintCondition(iter)) info.Print();
	OutputProgress(info);
      }

      if (stop) break;
    }

    if (MPI_IS_MASTER && !this->par.output.empty()) {
      W_big.save(this->par.output+"_W");
      H_big.save(this->par.output+"_H");
    }
  }

  template<typename MAT> void MultNMF<MAT>::Combine(mat& X, mat& X_big) {
    vector<mat> all_x;
    gather(this->world, X, all_x, 0);
    if (this->world.rank() == 0) PART<mat>::cat(all_x, X_big);
    broadcast(this->world, X_big, 0);
  }

  template<typename MAT> void MultNMF<MAT>::Initialize() {
    vector<int> r, c, sr, sc;
    vector<mat> w, h;
    gather(this->world, mr, r, 0);
    gather(this->world, nc, c, 0);
    if (MPI_IS_MASTER) {
      sr = r; sc = c;
      sr[0] = 0; sc[0] = 0;
      FORE(i, 1, MPI_SIZE) {
	sr[i] = sr[i-1]+r[i-1];
	sc[i] = sc[i-1]+c[i-1];
      }
      if (sr.back()+r.back() != this->par.m || sc.back()+c.back() != this->par.n)
	FATAL("Check dimension of submatrices")

      W_big = randu<mat>(this->par.m, this->par.k);
      PART<mat>::split(W_big, sr, w);

      H_big = randu<mat>(this->par.n, this->par.k);
      PART<mat>::split(H_big, sc, h);
    }
    PART<mat>::broadcast(this->world, w, W);
    PART<mat>::broadcast(this->world, h, H);
  }

  template<typename MAT> void MultNMF<MAT>::ParInitialize() {
    REQUIRE_PARAM("cs");
    int cs = vm["cs"].as<int>() * world.rank();
    srand(cs);
    W = randu<mat>(mr, par.k);
    H = randu<mat>(nc, par.k);
  }

  template<typename MAT> void MultNMF<MAT>::ComputeNormA(MAT& Arows) {
    double f = normAlocal = norm(Arows, "fro"); 
    f *= f;
    normA = 1;
    reduce(this->world, f, normA, std::plus<double>(), 0);
    normA = sqrt(normA);
  }

  template<typename MAT> void MultNMF<MAT>::ComputeNormE(MAT& Acols) {
    double f = normElocal = norm(Acols-W_big*trans(H),"fro"), s = 0;
    f *= f;
    normE = 1;
    reduce(this->world, f, normE, std::plus<double>(), 0);
    normE = sqrt(normE);
  }

  template<typename MAT> bool MultNMF<MAT>::StopCriteriaMet() {
    return relE < this->par.tol || fabs(relE-oldE) < this->par.tol
      || fabs(relE-oldE)/oldE < this->par.tol;
  }

  template<typename MAT> bool MultNMF<MAT>::PrintCondition(int iter) {
    return this->par.verbose == 2 || iter == 0 || iter == this->par.maxiter-1 || stop;
  }
  
  template<typename MAT> void MultNMF<MAT>::OutputProgress(Iteration& info) {
    if (!this->par.output.empty()) {
      ofstream f(this->par.output.c_str(), ios::app);
      info.Output(f);
    }
  }

}

